import re

def parse_log_file(file_path, skip):
    """从日志文件中提取性能指标数据"""
    elapsed_times = []
    tflops_throughputs = []
    token_throughputs = []
    
    # 正则表达式模式（优化版）
    pattern_iter = re.compile(r'iteration\s+(\d+)/')
    pattern_elapsed = re.compile(r'elapsed time per iteration \(ms\): ([\d\.]+)')
    pattern_tflops = re.compile(r'throughput \(tflops/sec/GPU\) : ([\d\.]+)')
    pattern_token = re.compile(r'throughput \(token/sec/GPU\) : ([\d\.]+)')
    
    try:
        with open(file_path, 'r') as file:
            for line in file:
                # 检查是否包含关键指标
                if "elapsed time per iteration" in line:
                    # 提取迭代编号
                    iter_match = pattern_iter.search(line)
                    if iter_match:
                        iter_num = int(iter_match.group(1))
                        # 从第二次迭代开始记录
                        if iter_num >= skip:
                            # 提取三项指标
                            elapsed_match = pattern_elapsed.search(line)
                            tflops_match = pattern_tflops.search(line)
                            token_match = pattern_token.search(line)
                            
                            if elapsed_match and tflops_match and token_match:
                                elapsed_times.append(float(elapsed_match.group(1)))
                                tflops_throughputs.append(float(tflops_match.group(1)))
                                token_throughputs.append(float(token_match.group(1)))
    
    except FileNotFoundError:
        print(f"错误：文件 '{file_path}' 不存在")
        return None
    except Exception as e:
        print(f"处理文件时出错：{str(e)}")
        return None
    
    return elapsed_times, tflops_throughputs, token_throughputs

def calculate_averages(metrics):
    """计算各项指标的平均值"""
    if not metrics or not all(metrics):
        print("未找到有效数据，请检查日志格式")
        return None
    
    elapsed, tflops, tokens = metrics
    avg_elapsed = sum(elapsed) / len(elapsed)
    avg_tflops = sum(tflops) / len(tflops)
    avg_tokens = sum(tokens) / len(tokens)
    
    return {
        "avg_elapsed_ms": round(avg_elapsed, 2),
        "avg_tflops": round(avg_tflops, 2),
        "avg_token_throughput": round(avg_tokens, 2),
        "iterations_counted": len(elapsed)
    }

# ===== 主程序 =====
if __name__ == "__main__":
    # 配置日志文件路径（修改为您的实际路径）
    # log_file = "/mnt/x10000/002266/output_fp8/logs/2025-08-02_01:37:27/tp1_pp8_dp16_mbs1_numbs_gbs1024_gpus128.15..log"
    log_file = "/mnt/x10000/002266/output_fp8/logs/fp8_960_8.8/tp1_pp8_dp120_mbs1_numbs_gbs19200_gpus960.119..log"
    skip = 203
    # 解析日志文件
    metrics = parse_log_file(log_file, skip)
    # print(metrics[1])
    if metrics:
        # 计算统计结果
        results = calculate_averages(metrics)
        
        # 打印结果表格
        print("="*55)
        print(f"{'性能指标统计':^55}")
        print("="*55)
        print(f"{'统计范围':<25} : 第{skip}次到第{results['iterations_counted']+skip-1}次迭代")
        print(f"{'平均迭代耗时':<25} : {results['avg_elapsed_ms']} ms")
        print(f"{'平均TFLOPS吞吐量':<25} : {results['avg_tflops']} tflops/sec/GPU")
        print(f"{'平均MFU':<25} : {results['avg_tflops']/458*100:.2f}%")
        print(f"{'平均Token吞吐量':<25} : {results['avg_token_throughput']} tokens/sec/GPU")
        print("="*55)


